from __future__ import annotations

import gymnasium as gym
import numpy as np

from driving_gym.data.data_provider import Snapshot
from driving_gym.environment.agent.obs.base_obs import BaseObs
from driving_gym.environment.scenario.carla_scenario import CarlaScenario
from driving_gym.misc.util import override
from driving_gym.simulation.adapter.carla.carla_adapter import CarlaAdapter
from driving_gym.simulation.adapter.carla.carla_config import CarlaActorConfig
from driving_gym.simulation.adapter.carla.planner.navigation.nav_planner import (
    get_latlon_ref,
    location_to_gps,
)
from driving_gym.simulation.adapter.carla.utils import carla_tranform


class WaypointObs(BaseObs):
    adapter: CarlaAdapter

    def __init__(self, config: dict, adapter: CarlaAdapter):
        super().__init__(config, adapter)
        self.gps = self.config.get("gps", False)
        self.num_waypoints = self.config.get("num_waypoints", 5)
        self.dist_to_actor = self.config.get("dist_threshold", 3)
        self.lon_ref = None
        self.lat_ref = None

    @override(BaseObs)
    def get_obs(self, snapshot: Snapshot, scenario: CarlaScenario) -> np.ndarray:
        """Return the following waypoint or goal location information

        Args:
            snapshot (Snapshot): One frame of the simulation data

        Returns:
            np.ndarray:
        """
        task = scenario.get_task(self.actor_id)
        goal_trans = CarlaActorConfig.parse_location(task["end"], self.adapter.map)
        goal_loc = goal_trans.location
        vehicle_loc = self.adapter.get_transform(self.actor_id).location

        waypoints = []
        if scenario.enable_planner:
            planner = scenario.planners[self.actor_id]
            wpts = planner.get_nearest_waypoints(20)
            for wpt in wpts:
                loc = wpt.transform.location
                if vehicle_loc.distance(loc) > self.dist_to_actor:
                    waypoints.append(loc)
                    if len(waypoints) >= self.num_waypoints - 1:
                        break
        else:
            wpt = self.adapter.map.get_waypoint(vehicle_loc)
            while len(waypoints) < self.num_waypoints - 1:
                next_wpts = wpt.next()
                if not next_wpts:
                    break

                wpt = next_wpts[0]
                loc = wpt.transform.location
                if vehicle_loc.distance(loc) > self.dist_to_actor:
                    waypoints.append(loc)

        while len(waypoints) < self.num_waypoints:
            waypoints.append(goal_loc)

        if self.gps:
            if self.lon_ref is None:
                self.lat_ref, self.lon_ref = get_latlon_ref(self.adapter.map)

            coords = []
            for location in waypoints:
                gps_coord = location_to_gps(self.lat_ref, self.lon_ref, location)
                coords.extend([gps_coord["lat"], gps_coord["lon"], gps_coord["z"]])
            return np.array(coords, dtype=np.float32)
        else:
            coords = []
            for location in waypoints:
                coords.extend([location.x, location.y, location.z])
            return np.array(coords, dtype=np.float32)

    @override(BaseObs)
    def get_observation_space(self) -> gym.spaces.Box:
        return gym.spaces.Box(
            low=-np.inf, high=np.inf, shape=(3 * self.num_waypoints,), dtype=np.float32
        )
